# -*- coding: utf-8 -*-
# created on 2016/9/5


from sympy import *
from mathsolver.functions.base import *


class TrigSimplify(BaseFunction):
    def __init__(self):
        super(TrigSimplify, self).__init__()
        self.trlist = []
        self.expr = None

    @staticmethod
    def is_sin(eq):
        return isinstance(eq, sin)

    def is_abs_sin(self, eq):
        if self.is_sin(eq):
            return True
        elif eq.is_Mul and eq.args[0] == -1 and self.is_sin(eq.args[1]):
            return True
        else:
            return False

    @staticmethod
    def is_cos(eq):
        return isinstance(eq, cos)

    def is_sin_cos(self, eq):
        return self.is_sin(eq) or self.is_cos(eq)

    def is_sincos(self, eq1, eq2):
        return self.is_sin(eq1) and self.is_cos(eq2)

    def is_abs_cos(self, eq):
        if self.is_cos(eq):
            return True
        elif eq.is_Mul and eq.args[0] == -1 and self.is_cos(eq.args[1]):
            return True
        else:
            return False

    def is_abs_sin_cos(self, eq):
        return self.is_abs_sin(eq) or self.is_abs_cos(eq)

    def is_abs_sincos(self, eq1, eq2):
        return self.is_abs_sin(eq1) and self.is_abs_cos(eq2)

    @staticmethod
    def is_tan(eq):
        return isinstance(eq, tan)

    @staticmethod
    def is_cot(eq):
        return isinstance(eq, cot)

    @staticmethod
    def is_spe_dig(eq):
        """
        判定是否为特殊的三角函数值
        :param eq: 表达式
        :return: True or False
        """
        eq = fraction(eq)
        if ((eq[0] == sqrt(3) and (eq[1] == 3 or eq[1] == 2 or eq[1] == 1)) or
                (eq[0] == sqrt(2) and (eq[1] == 2 or eq[1] == 1)) or (eq[0] == 1 and eq[1] == 2)):
            return True
        else:
            return False

    def find(self, eq):
        """
        解决递归函数初始化问题,然后调用递归
        :param eq: 表达式
        :return: self
        """
        self.trlist = [[] for _ in range(15)]
        self.find_pattern(eq)
        return self

    def find_pattern(self, eq):
        """
        递归遍历表达式,将符合所有化简条件的项通过类内list传出来
        :param eq: 表达式
        :return: self
        """
        if isinstance(eq, (Symbol, Number)) or eq == sympify('pi'):
            return self
        if len(eq.args) > 0:
            if eq.is_Pow:
                # sin(x)**2  cos(x)**4
                if eq.args[1].is_even:
                    if self.is_sin(eq.args[0]):
                        self.trlist[5].append(eq)
                    elif self.is_cos(eq.args[0]):
                        self.trlist[6].append(eq)
                        self.trlist[7].append(eq)

            elif eq.is_Mul:
                for i in range(len(eq.args) - 1):
                    # sin(x)*cos(x)   cos(y)*sin(x)  tan(x)*tan(y)  cot(x)*cot(y)
                    if self.is_sin_cos(eq.args[i]) and self.is_sin_cos(eq.args[i + 1]):
                        if eq.args[i].args[0] == eq.args[i + 1].args[0]:
                            self.trlist[11].append([eq.args[i] * eq.args[i + 1], sin(2 * eq.args[i].args[0]) * 1 / 2])
                        else:
                            self.trlist[8].append(eq.args[i] * eq.args[i + 1])
                    elif (self.is_tan(eq.args[i]) and self.is_tan(eq.args[i + 1])) or \
                            (self.is_cot(eq.args[i]) and self.is_cot(eq.args[i + 1])):
                        self.trlist[13].append(eq.args[i] * eq.args[i + 1])

            elif eq.is_Add:
                for i in range(len(eq.args) - 1):
                    for j in range(i + 1, len(eq.args)):
                        v1 = eq.args[i]
                        v2 = eq.args[j]
                        # cos(x)/2 + sin(y)/2   sin(x)-cos(y)
                        if (self.is_abs_sin(v1) and self.is_abs_sin(v2)) or (
                                self.is_abs_cos(v1) and self.is_abs_cos(v2)):
                            self.trlist[9].append(v1 + v2)
                        elif v1.is_Mul and v2.is_Mul:
                            t1 = v1 / v1.args[len(v1.args) - 1]
                            t2 = v2 / v2.args[len(v2.args) - 1]
                            v1 = v1.args[len(v1.args) - 1]
                            v2 = v2.args[len(v2.args) - 1]
                            if (t1 == t2 or t1 == -t2) and \
                                    ((self.is_abs_sin(v1) and self.is_abs_sin(v2)) or (
                                            self.is_abs_cos(v1) and self.is_abs_cos(v2))):
                                self.trlist[9].append(eq.args[i] + eq.args[j])
                        elif v1.is_Mul and self.is_abs_sin_cos(v2):
                            v1 = v1.args[len(v1.args) - 1]
                        elif self.is_abs_sin_cos(v1) and v2.is_Mul:
                            v2 = v2.args[len(v2.args) - 1]
                        # sin(x)+cos(x)  sqrt(3)*sin(x)+1/2*cos(x)
                        if self.is_sincos(v1, v2) and v1.args[0] == v2.args[0]:
                            self.trlist[0].append([eq.args[i], eq.args[j], v1.args[0]])
                        elif self.is_sincos(v2, v1) and v1.args[0] == v2.args[0]:
                            self.trlist[0].append([eq.args[j], eq.args[i], v1.args[0]])

            # sin(x)   cos(x)
            elif self.is_sin_cos(eq):
                self.trlist[14].append(eq)
                if len(eq.args[0].free_symbols) > 0:
                    if eq.args[0].is_Mul and eq.args[0].args[0].is_even:
                        self.trlist[11].append([eq])
                    elif not eq.args[0].is_Symbol:
                        self.trlist[10].append(eq)
                else:
                    self.trlist[3].append(eq)

            # sec / csc
            elif isinstance(eq, sec) or isinstance(eq, csc):
                self.trlist[1].append(eq)

            # cot
            elif self.is_cot(eq):
                self.trlist[2].append(eq)

            # tan
            elif self.is_tan(eq):
                self.trlist[2].append(eq)
                if len(eq.args[0].free_symbols) > 0:
                    self.trlist[14].append(eq)
                    if eq.args[0].is_Mul and eq.args[0].args[0].is_even:
                        self.trlist[11].append(eq)
                    elif not eq.args[0].is_Symbol:
                        self.trlist[12].append(eq)

            for ag in eq.args:
                if self.is_spe_dig(ag):
                    self.trlist[4].append(ag)
                else:
                    self.find_pattern(ag)

            return self

    @staticmethod
    def lenexp(eq):
        """
        返回eq的长度. 长度定义:为多项式的项数,如果结果为纯数字,则认为化简到最优情况,l=0.
        例外: (sqrt(3)/3).is_Number 不为 0  但其l应该为0 ; (1+sqrt(3)) l也应该为0 ; sin(pi/5)*sin(pi/7)*sin(pi/9)  l=1
        :param eq: 表达式
        :return: l 表达式长度
        """
        if eq.is_Add:
            cnt = 0
            for tmp in eq.args:
                if isinstance(tmp, (sin, cos, tan, cot)) or len(tmp.free_symbols) > 0:
                    cnt += 1
            ll = cnt
        elif eq.is_Mul:
            ll = 0
            for tmp in eq.args:
                if isinstance(tmp, (sin, cos, tan, cot)):
                    ll = 1
        elif len(eq.free_symbols) == 0:
            ll = 0
        else:
            ll = 1
        return ll

    # 多项式化简
    @staticmethod
    def TR(item):
        """
        先调用一次factor针对分式型表达式进行公因式化简,再调用依次expand将多项式展开成乘累加形式
        :param item: 表达式
        :return: exp: 化简后表达式
        """
        ep = expand(factor(item))
        return ep

    # 辅助角公式
    def TR0(self, item):
        if len(self.trlist[0]) > 0:
            for x1, x2, sym in self.trlist[0]:
                t1 = fraction(x1)
                t2 = fraction(x2)
                v1 = t1[0].as_coefficient(sin(sym)) / t1[1]
                v2 = t2[0].as_coefficient(cos(sym)) / t2[1]
                v1, v2 = sqrt(v1 ** 2 + v2 ** 2), atan(v2 / v1)
                if not isinstance(v2, atan):
                    eq = x1 + x2
                    vv = sin(sym + v2) * v1
                    item = item.subs(eq, vv)
        return item

    # 正割、余割转换为正弦 、余弦的公式
    def TR1(self, item):
        for eq in self.trlist[1]:
            v = eq.args[0]
            if isinstance(arg, sec):
                vv = 1 / cos(v)
            else:
                vv = 1 / sin(v)
            item = item.subs(arg, vv)
        return item

    # 正切、余切转换为正弦 、余弦的公式
    def TR2(self, item):
        for eq in self.trlist[2]:
            v = eq.args[0]
            if isinstance(eq, tan):
                vv = sin(v) / cos(v)
            else:
                vv = cos(v) / sin(v)
            item = item.subs(eq, vv)
        return item

    # 诱导公式  sympy自动转换
    def TR3(self):
        return

    # 特殊角的三角函数值  sympy自动求值
    def TR4(self):
        return

    # 正弦平方替换
    def TR5(self, item):
        varlist = [(1 - cos(eq.args[0].args[0]) ** 2) ** (eq.args[1] / 2) for eq in self.trlist[5]]
        dic = zip(self.trlist[5], varlist)
        item = item.subs(dic)
        return item

    # 余弦平方替换
    def TR6(self, item):
        varlist = [(1 - sin(eq.args[0].args[0]) ** 2) ** (eq.args[1] / 2) for eq in self.trlist[6]]
        dic = zip(self.trlist[6], varlist)
        item = item.subs(dic)
        return item

    # 余弦平方降次
    def TR7(self, item):
        for eq in self.trlist[7]:
            v = eq.args[0].args[0]
            vv = (1 + cos(2 * v)) / 2
            item = item.subs(eq, vv)
        return item

    # 积化和差
    def TR8(self, item):
        for eq in self.trlist[8]:
            v1 = eq.args[0].args[0]
            v2 = eq.args[1].args[0]
            if self.is_sin(eq.args[0]) and self.is_cos(eq.args[1]):
                vv = (sin(v1 + v2) + sin(v1 - v2)) / 2
            elif self.is_cos(eq.args[0]) and self.is_sin(eq.args[1]):
                vv = (sin(v1 + v2) - sin(v1 - v2)) / 2
            elif self.is_cos(eq.args[0]) and self.is_cos(eq.args[1]):
                vv = (cos(v1 + v2) + cos(v1 - v2)) / 2
            elif self.is_sin(eq.args[0]) and self.is_sin(eq.args[1]):
                vv = - (cos(v1 + v2) - cos(v1 - v2)) / 2
            vv = expand(factor(vv))
            item = item.subs(eq, vv)
        return item

    # 和差化积
    def TR9(self, item):
        for eq in self.trlist[9]:
            t1, t2 = eq.args[0], eq.args[1]
            t = 1
            if t1.is_Mul and t2.is_Mul:
                t = t1 / t1.args[len(t1.args) - 1]
                t1 = t1 / t
                t2 = t2 / t
            if t1.is_Mul and not t2.is_Mul:
                v1 = t2.args[0]
                v2 = t1.args[1].args[0]
                if self.is_sin(t2):
                    vv = 2 * t * cos((v1 + v2) / 2) * sin((v1 - v2) / 2)
                else:
                    vv = -2 * t * sin((v1 + v2) / 2) * sin((v1 - v2) / 2)
            elif not t1.is_Mul and t2.is_Mul:
                v1 = t1.args[0]
                v2 = t2.args[1].args[0]
                if self.is_sin(t1):
                    vv = 2 * t * cos((v1 + v2) / 2) * sin((v1 - v2) / 2)
                else:
                    vv = -2 * t * sin((v1 + v2) / 2) * sin((v1 - v2) / 2)
            else:
                v1 = t1.args[0]
                v2 = t2.args[0]
                if self.is_sin(t1):
                    vv = 2 * t * sin((v1 + v2) / 2) * cos((v1 - v2) / 2)
                else:
                    vv = 2 * t * cos((v1 + v2) / 2) * cos((v1 - v2) / 2)
            vv = expand(factor(vv))
            item = item.subs(eq, vv)
            return item
        return exp

    # 两角和与差公式
    def TR10(self, item):
        for eq in self.trlist[10]:
            v = eq.args[0]
            if len(v.args) == 2:
                if v.is_Add:
                    v1, v2 = v.args[0], v.args[1]
                elif v.is_Mul and v.args[0].is_Number and v.args[0].is_odd:  # tan(3*x)
                    v1, v2 = (v.args[0] + 1) / 2 * v.args[1], (v.args[0] - 1) / 2 * v.args[1]
                if self.is_sin(eq):
                    vv = sin(v1) * cos(v2) + cos(v1) * sin(v2)
                elif self.is_cos(eq):
                    vv = cos(v1) * cos(v2) - sin(v1) * sin(v2)
                item = item.subs(eq, vv)
        return item

    # 倍角公式
    def TR11(self, item):
        for eq in self.trlist[11]:
            if len(eq) == 1:
                v = eq[0].args[0]
                if self.is_sin(eq[0]):
                    vv = 2 * sin(v / 2) * cos(v / 2)
                elif self.is_tan(eq[0]):
                    vv = 2 * tan(v / 2) / (1 - tan(v / 2) ** 2)
                else:
                    vv = cos(v / 2) ** 2 - sin(v / 2) ** 2
                item = item.subs(eq[0], vv)
            elif len(eq) == 2:
                item = item.subs(eq[0], eq[1])
        return item

    # 两角和与差的正切
    def TR12(self, item):  # tan(a+b)  tan(a-b)
        for eq in self.trlist[12]:
            v = eq.args[0]
            if len(v.args) == 2:  # 不超过两项, 比如不处理 tan(x+y+z) 或 tan(3*x*y)
                if v.is_Add:  # add不能超过两项  比如tan(x+y)
                    v1, v2 = v.args[0], v.args[1]
                    vv = (tan(v1) + tan(v2)) / (1 - tan(v1) * tan(v2))
                elif v.is_Mul and v.args[0].is_Number and v.args[0].is_odd:  # tan(3*x)
                    v1, v2 = (v.args[0] + 1) / 2 * v.args[1], (v.args[0] - 1) / 2 * v.args[1]
                    vv = (tan(v1) + tan(v2)) / (1 - tan(v1) * tan(v2))
                item = item.subs(eq, vv)
        return item

    # 两角正余切的积
    def TR13(self, item):
        for eq in self.trlist[13]:
            v1 = eq.args[0].args[0]
            v2 = eq.args[1].args[0]
            if self.is_tan(eq.args[0]) and self.is_tan(eq.args[1]):
                vv = 1 - (tan(v1) + tan(v2)) * cot(v1 + v2)
            else:
                vv = 1 + (cot(v1) + cot(v2)) * cot(v1 + v2)
            vv = expand(factor(vv))
            item = item.subs(eq, vv)
        return item

    # 万能公式
    def TR14(self, item):
        for eq in self.trlist[14]:
            v = eq.args[0]
            if self.is_sin(eq):
                vv = 2 * tan(v / 2) / (1 + tan(v / 2) ** 2)
            elif self.is_cos(eq):
                vv = (1 + tan(v / 2) ** 2) / (1 + tan(v / 2) ** 2)
            else:
                vv = 2 * tan(v / 2) / (1 - tan(v / 2) ** 2)
            item = item.subs(eq, vv)
        return item

    # TR5 | TR6
    def CTR1(self, item):
        expr1 = self.TR(self.find(item).TR5(item))
        expr2 = self.TR(self.find(item).TR6(item))
        label = ''
        if self.lenexp(expr1) < self.lenexp(item) and self.lenexp(expr1) <= self.lenexp(expr2):
            item = expr1
            label = '正弦平方替换'
        elif self.lenexp(expr2) < self.lenexp(item) and self.lenexp(expr2) <= self.lenexp(expr1):
            item = expr2
            label = '余弦平方替换'
        return item, label

    # TR11 + (TR5 | TR6)
    def CTR2(self, item):
        expr3 = self.TR(item)
        expr3 = self.TR(self.find(expr3).TR11(expr3))
        self.find(expr3)
        expr1 = self.TR(self.TR5(expr3))
        expr2 = self.TR(self.TR6(expr3))
        label = '(倍角公式)'
        if self.lenexp(expr1) < self.lenexp(expr3) and self.lenexp(expr1) <= self.lenexp(expr2):
            item = expr1
            label = '(倍角公式 正弦平方转换)'
        elif self.lenexp(expr2) < self.lenexp(expr3) and self.lenexp(expr2) <= self.lenexp(expr1):
            item = expr2
            label = '(倍角公式 余弦平方转换)'
        else:
            item = expr3
        return item, label

    # TR8 + TR10 -1
    def CTR3(self, item):
        expr1 = self.TR(item)
        self.find(expr1)
        while len(self.trlist[8]) > 0:
            expr1 = self.TR(self.find(expr1).TR8(expr1))
            self.find(expr1)
        expr2 = self.TR(expr1)
        expr2 = self.find(expr2).TR0(expr2)
        label = ''
        if self.lenexp(expr2) < self.lenexp(item):
            item = expr2
            label = '(积化和差 辅助角公式)'
        elif self.lenexp(expr1) < self.lenexp(item):
            item = expr1
            label = '(积化和差)'
        return item, label

    # TR4 + TR10 -1 = TR0 辅助角公式
    def CTR4(self, item):
        expr1 = self.find(item).TR0(item)
        label = ''
        if self.lenexp(expr1) < self.lenexp(item):
            item = expr1
            label = '(辅助角公式)'
        return item, label

    def TR_RuleList(self, item, rule_list):
        """
        根据ruleList依次对expr调用对应的化简序列, 化简方法集合为{TR0-14,CTR1-4}
        :param item: 待化简表达式
        :param rule_list: 化简规则列表
        :return: expr 化简后的表达式
        """
        for lii in rule_list:
            item = self.TR(item)
            self.find(item)
            if len(self.trlist[lii]) > 0:
                if lii == 0:  # CTR4
                    item, label = self.CTR4(item)
                elif lii == 1:
                    item = self.TR(self.TR1(item))
                    label = '(正/余割转正/余弦)'
                elif lii == 2:
                    item = self.TR(self.TR2(item))
                    label = '(正/余切转正/余弦)'
                elif lii == 5:
                    item = self.TR(self.TR5(item))
                    label = '(正弦平方替换)'
                elif lii == 6:  # CTR1
                    item, label = self.CTR1(item)
                elif lii == 7:
                    item = self.TR(self.TR7(item))
                    label = '(余弦平方降次)'
                elif lii == 8:  # CTR3
                    item, label = self.CTR3(item)
                elif lii == 9:
                    # 存在需要进行多个和差化积化简的情况,所以循环调用,知道所有可和差化积的项全部化简完为止
                    while len(self.trlist[lii]) > 0:
                        item = self.TR(self.TR9(item))
                        self.find(item)
                    label = '(和差化积)'
                elif lii == 10:
                    item = self.TR(self.TR10(item))
                    label = '(两角和与差公式)'
                elif lii == 11:  # CTR2
                    item, label = self.CTR2(item)
                elif lii == 12:
                    item = self.TR(self.TR12(item))
                    label = '(两角和与差的正切)'
                elif lii == 13:
                    item = self.TR(self.TR13(item))
                    label = '(两角正余切的积)'
                elif lii == 14:
                    item = self.TR(self.TR14(item))
                    label = '(万能公式)'
                if label:
                    print('\t=', item, label)
        return item

    def printList(self):
        """
        打印出self.trlist,用于逻辑debug,确认化简项提取正确
        :return:
        """
        for i in range(15):
            print('TR', i, ':', self.trlist[i])

    def doList(self, item):
        """
        利用一套可以化简常见三角函数的化简规则进行三角函数化简
        1 调用规则1,2,12,13将表达式化简为只含有sin cos项
        2 调用10,11将和角和倍角展开化为单角的形式,尽量将不同三角函数转化为同一三角函数
        3 三角函数表达式的降次,合并: sin偶次方化成cos偶次方,对cos偶次方降次,和差化积并对表达式进行多项式合并/化简
        4 三角函数表达式的收敛,聚合: 用和差化积/多角和与差公式/倍角公式/组合规则(CTR1-4)
        :param item: 待化简表达式
        :return: expr  化简后表达式
        """
        rl1 = [12, 13]
        # 10-两角和与差 11-倍角 5-正弦转余弦 7-余弦平方降次 9-和差化积 11-倍角公式 8-积化和差/辅助角公式 6-正余弦转换 0-辅助角公式
        rl2 = [10, 11] + [5, 7, 9, 11] + [8, 6, 9, 11, 9, 9, 0]

        self.find(item)
        if len(self.trlist[1]) > 0:
            item = self.TR1(item)
            print('\t=', item, '(正/余割转正/余弦)')
        self.find(item)
        if len(self.trlist[2]) > 0:
            expr1 = self.TR_RuleList(item, rl1)
            if self.lenexp(expr1) < self.lenexp(item):
                item = expr1
        self.find(item)
        if len(self.trlist[2]) > 0:
            item = self.TR(self.TR2(item))
            print('\t=', item, '(正/余切转正/余弦)')
        self.find(item)
        if len(self.trlist[14]) > 0:
            item = self.TR_RuleList(item, rl2)
        return item

    def tryList(self, item):
        """
        目前调用一套给定的规则序列必然不能解决所有化简问题,应该使用一套化简思想,然后让程序在这个解题化简思路情况下自行去查找解决方案更优.
        但对于化简结果的判定函数就显得很重要,结果的价值判定函数好与坏将影响化简的好坏.
        //尚未完成
        :param item: 表达式
        :return:
        """
        rl1 = [1, 12, 13, 10, 11]
        item = self.TR_RuleList(item, rl1)
        self.find(item)
        expr1 = self.TR5(item)
        if expr1 != item:
            if Abs(self.lenexp(expr1) - self.lenexp(item)) <= 3:
                item = expr1
                print('\t=', item, '(正/余弦平方替换)')
            self.find(item)
            expr1 = self.TR0(self.TR7(item))
            if Abs(self.lenexp(expr1) - self.lenexp(item)) <= 3:
                item = expr1
                print('\t=', item, '(余弦平方降次)')
        self.find(item)

    def solver(self, *args):
        f1, f2 = args[0].sympify()
        self.expr = f1 - f2
        print('=======================')
        print('原式 =', self.expr)

        # 由于套路固定,会将多角/倍角/tan/cot等全部进行转换然后化简的套路,
        # 对于不需要化简而是考虑整体代换化简的题目往往不会得到好的化简结果
        self.doList(self.expr)
        print('trigsimp结果:', trigsimp(self.expr))
        return self


if __name__ == '__main__':
    TrigSimplify().solver(BaseEq(['tan(x) - cot(x)', 0]))
    pass
